import argparse
import cv2
import io
from flask import Flask, render_template, Response

import pyrealsense2 as rs
import  numpy as np



class RealSense:
    def __init__(self):
        self.pipeline = rs.pipeline()
        config = rs.config()
        H, W = 480, 640
        self.H, self.W = H, W
        # config.enable_stream(rs.stream.depth, W, H, rs.format.z16, 30)
        config.enable_stream(rs.stream.color, W, H, rs.format.bgr8, 30)
        cfg = self.pipeline.start(config)
    
    def read(self):
        while True:
            frames = self.pipeline.wait_for_frames()
            # depth_frame = frames.get_depth_frame()
            color_frame = frames.get_color_frame()
            if not color_frame:
                continue
            color_image = np.asanyarray(color_frame.get_data())
            return 1, color_image


class Cam:
    def __init__(self, source):
        self.cap = cv2.VideoCapture(source)
        if self.cap.isOpened():
            rval, frame = self.cap.read()
    
    def read(self):
        while True:
            rval, frame = self.cap.read()
            if rval:
                yield 1, frame


app = Flask(__name__)
cam = None


@app.route('/')
def index():  
    return render_template('index.html')


def gen():
    while True:
        ret, frame = cam.read()
        # yield frame
        encode_return_code, image_buffer = cv2.imencode('.jpg', frame)
        io_buf = io.BytesIO(image_buffer)
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + io_buf.read() + b'\r\n')


@app.route('/video_feed')
def video_feed():
    return Response(gen(), mimetype='multipart/x-mixed-replace; boundary=frame')


def parse_argument():
    parser = argparse.ArgumentParser(description="video stream server")
    parser.add_argument("-s", "--source", type=str)
    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = parse_argument()
    source = args.source
    if "realsense" in source:
        cam = RealSense()
    else:
        cam = Cam(source)
    
    # when runing with realsense set debug is False. I don't know why.
    app.run(host='0.0.0.0', debug=False, threaded=True)
